import pandas as pd
import numpy as np
import os

print("收益回测开始~")

# 如果当前路径没有"_3.策略表现"文件夹，则创建该文件夹
if not os.path.exists("_3.策略表现"):
    os.mkdir("_3.策略表现")


def get_stk_return(stk_code, col):
    ps_i = (df_stk["证券代码"] == stk_code)
    sr = df_stk[col][ps_i].values[0] / 100  # 处理涨跌幅月的单位
    stk_return = sr if not np.isnan(sr) else 0  # 处理缺失值的问题

    return stk_return


def get_sty_return(df_sty, col):
    sty_return = 0
    for i in range(len(df_sty)):
        sr = get_stk_return(df_sty["证券代码"][i], col)
        wt = df_sty["金额比例"][i]
        sty_return += (1 + sr) * wt

    return sty_return - 1


df_stk = pd.read_excel("源数据\\1.月度涨跌幅.xlsx")  # 个股收益率
df_idx = pd.read_excel("源数据\\4.东方财富全A.xlsx")  # 指数收益率

# df_策略表现
df_rst = pd.DataFrame(columns=["序号", "文件名", "组合收益率", "指数收益率", "是否战胜市场"])

# --------------------------------------------------------------------------计算回测收益率
"""
月初买入，月底(下个月初)卖出，用这个区间的收益率做回测
"2025-04-01(组合).xlsx"这个股票组合的收益率得等5月初才有
"""
files = ["2025-01-01(组合).xlsx", "2025-02-01(组合).xlsx", "2025-03-01(组合).xlsx"]
cols = ["2025-01-27", "2025-02-28", "2025-03-31"]

for i in range(len(files)):
    file = files[i]
    col = cols[i]

    df_sty = pd.read_excel("_2.月度组合\\" + file)

    idx_return = df_idx[col][0] / 100
    sty_return = get_sty_return(df_sty, col)

    df_rst.loc[i, "组合收益率"] = sty_return
    df_rst.loc[i, "指数收益率"] = idx_return
    df_rst.loc[i, "是否战胜市场"] = int(sty_return > idx_return)

df_rst["序号"] = range(1, len(files) + 1)
df_rst["文件名"] = files

df_rst.to_excel("_3.策略表现\\策略表现.xlsx", index=False)

print("收益回测完成！")
